
import sys
from config import *
import coveragefromsummary

def lookupGeneInterval(gene, filename=KNOWN_GENE_TABLE):
    for line in open(filename):
        if line[0]=="#": continue
        cols = line.strip().split("\t")
        if gene != cols[0]: continue
        return (cols[1], cols[2], long(cols[3]), long(cols[4]))
    print >> sys.stderr, "# Couldn't find gene ", gene
    return None

def trackName(queryId, samplefile, samplename):
    if samplefile != samplename:
        return queryId + "--" + samplename
    else:
        return queryId + "--" +samplefile.split("/")[-1]
    
def generateGffCoverage(requests, samplefiles, filename=KNOWN_GENE_TABLE):

    queries = {}
    for q in requests:
        if len(q)==3:
            # query of form ("chr21", 1, 1000)
            (chr, spos, epos) = (q[0], long(q[1]), long(q[2]))
            queryId = "%s:%s-%s" % (chr, spos, epos)
            queries[queryId] = (chr, '?', spos, epos)
        elif len(q)==2:
            (gene, utrLength) = (q[0], long(q[1]))
            (chr, strand, spos, epos) = lookupGeneInterval(gene)            
            queries[gene] = (chr, '?', spos - utrLength, epos + utrLength)
        else:
            raise Exception("bad coverage request: " + str(q))
    
    for queryId in queries.keys():
        print >> sys.stderr, "Query '%s' @ %s" % (queryId, queries[queryId])

    printedHeader = False
    for queryId in queries.keys():
        print >> sys.stderr, "Processing query '%s' @ %s ..." % (queryId, queries[queryId])
        (chr, strand, spos, epos) = queries[queryId]
        if not printedHeader:
            printedHeader = True
            print "browser position %s:%ld-%ld" % (chr, spos, epos)
        for (samplefile, samplename) in samplefiles:
            anyCoverage = False
            print "track name=\"%s\" visibility=full" % trackName(queryId, samplefile, samplename)
            # cov = calculateLocalCoverage(queries[queryId], samplefile)
            coverage = coveragefromsummary.calculateCoverageGenerator(queries[queryId], samplefile)
            prev = 0.0
            prevAt = spos
            featureId = "ID=coverage-%s-%s;Sample=%s;Gene=%s" % (queryId, samplename, samplename, queryId)
            # featureId = ""
            # featureSource = "cov"
            featureSource = "cov%s" % samplename
            # featureType = "cov"
            featureType = "cov%s" % samplename
            for (x, cov) in coverage:
                if cov != prev:
                    if prev != 0.0:
                        anyCoverage = True
                    print "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s" % (chr, featureSource, featureType, prevAt, x -1, int(prev), ".", ".", featureId)
                    prevAt = x
                    prev = cov
            # if prev != 0.0 or not anyCoverage:
            print "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s" % (chr, featureSource, featureType, prevAt, epos+1 -1, int(prev), ".", ".", featureId)

def printUsage(args, usage, minArgs=2):
  if len(args) < minArgs or args[1].startswith("-h") or args[1].startswith("--h"):
    print usage
    sys.exit(1)

usageString = """
Usage: python %s QUERY SUMMARY_FILE1 SUMMARY_FILE2 ...,
where QUERY is either [-UTR utrLength] GENE or -Q CHR MIN_POS MAX_POS.

Generates a coverage map for the specified gene or interval, in GFF format,
from each of the specified summary files.  (If -UTR is specified, extend
the gene on either side by the given UTR length.)
"""

def getSampleName(samplefile):
    filename = samplefile.split("/")[-1]
    ignoreBits = ["illumina", "burge", "txt", "snp-results", "bowtie_sorted_summary", "INDEX"]
    bits = filter(lambda(x):(x not in ignoreBits), filename.split("."))
    return ".".join(bits)

if __name__=="__main__":
    printUsage(sys.argv, usageString % sys.argv[0], 3)
    if sys.argv[1]=='-Q':
        chr = sys.argv[2]
        spos = long(sys.argv[3])
        epos = long(sys.argv[4])
        query = (chr, spos, epos)
        samples = sys.argv[5:]
    elif sys.argv[1]=='-UTR':
        utrLength = int(sys.argv[2])
        gene = sys.argv[3]
        query = (gene, utrLength)
        samples = sys.argv[4:]
    else:
        gene = sys.argv[1]
        query = (gene, 0)
        samples = sys.argv[2:]
    generateGffCoverage([query], [(x,getSampleName(x)) for x in samples])
